import pandas as pd
from sqlalchemy import create_engine
from sqlalchemy.engine import URL
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Hive连接配置
HIVE_CONFIG = {
    'host': '192.168.0.12',
    'port': 10000,
    'database': 'default',
    'username': 'hiveuser',
    'password': 'hivepassword',
    'auth': 'LDAP'  # 根据实际情况修改认证方式
}


def get_hive_connection():
    """创建Hive连接"""
    connection_url = URL.create(
        drivername="hive",
        username=HIVE_CONFIG['username'],
        password=HIVE_CONFIG['password'],
        host=HIVE_CONFIG['host'],
        port=HIVE_CONFIG['port'],
        database=HIVE_CONFIG['database'],
        query={"auth": HIVE_CONFIG['auth']}
    )
    engine = create_engine(connection_url)
    return engine.connect()


def get_databases(connection):
    """获取所有数据库名"""
    query = "SHOW DATABASES"
    result = connection.execute(query)
    return [row[0] for row in result]


def get_tables(connection, database):
    """获取指定数据库中的所有表"""
    query = f"SHOW TABLES IN {database}"
    result = connection.execute(query)
    return [row[0] for row in result]


def get_table_comment(connection, database, table):
    """获取表注释"""
    try:
        query = f"DESCRIBE FORMATTED {database}.{table}"
        result = connection.execute(query)

        for row in result:
            if row[0] == 'Comment':
                return row[1] if row[1] else "无注释"
        return "无注释"
    except Exception as e:
        logger.error(f"获取表 {database}.{table} 注释失败: {e}")
        return "获取注释失败"


def get_column_count(connection, database, table):
    """获取表的字段数量"""
    try:
        query = f"DESCRIBE {database}.{table}"
        result = connection.execute(query)

        # 过滤掉分区字段和元数据行
        columns = [row for row in result if not row[0].startswith('#')]
        return len(columns)
    except Exception as e:
        logger.error(f"获取表 {database}.{table} 字段数量失败: {e}")
        return -1


def get_row_count(connection, database, table):
    """获取表的行数"""
    try:
        # 注意：COUNT(*) 可能对大表性能较差，考虑使用近似计数
        query = f"SELECT COUNT(*) FROM {database}.{table}"
        result = connection.execute(query)
        return result.fetchone()[0]
    except Exception as e:
        logger.error(f"获取表 {database}.{table} 行数失败: {e}")
        return -1


def export_to_excel(metadata_df, output_file):
    """导出元数据到Excel"""
    try:
        metadata_df.to_excel(output_file, index=False, engine='openpyxl')
        logger.info(f"元数据已成功导出到 {output_file}")
    except Exception as e:
        logger.error(f"导出Excel失败: {e}")


def main():
    """主函数"""
    output_file = "hive_metadata_export.xlsx"
    metadata = []

    try:
        with get_hive_connection() as connection:
            # 获取所有数据库
            databases = get_databases(connection)
            logger.info(f"找到 {len(databases)} 个数据库")

            # 遍历每个数据库
            for database in databases:
                # 获取数据库中的所有表
                tables = get_tables(connection, database)
                logger.info(f"数据库 {database} 包含 {len(tables)} 个表")

                # 遍历每个表
                for table in tables:
                    # 获取表的详细信息
                    comment = get_table_comment(connection, database, table)
                    column_count = get_column_count(connection, database, table)
                    row_count = get_row_count(connection, database, table)

                    # 收集元数据
                    metadata.append({
                        '数据库名': database,
                        '表名': table,
                        '表注释': comment,
                        '字段个数': column_count,
                        '行数': row_count
                    })

        # 转换为DataFrame并导出到Excel
        if metadata:
            metadata_df = pd.DataFrame(metadata)
            export_to_excel(metadata_df, output_file)
        else:
            logger.warning("未找到任何表元数据")

    except Exception as e:
        logger.error(f"程序执行失败: {e}")


if __name__ == "__main__":
    main()